from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.edge.options import Options
from selenium.webdriver.edge.service import Service as EdgeService
from webdriver_manager.microsoft import EdgeChromiumDriverManager
from selenium.common.exceptions import WebDriverException

service = EdgeService(executable_path=EdgeChromiumDriverManager().install())

class WebDriver:
    def __init__(self):
        options = Options()

        options.add_argument('--headless=new')#无头浏览器
        options.add_argument('--window-size=1920,1080')#设置浏览器尺寸
        options.add_argument('--ignore-certificate-errors')#https证书问题
        self.driver = driver = webdriver.Edge(service=service,options=options)
    def set_cookie(self, cookie):#支持字符串/列表/字典方式设置cookie
        match cookie:
            case int() | float():
                return False
            case str():#字符串
                self.driver.execute_script(
                    'document.cookie="' + cookie + '";'
                )
            case list():
                for thisCookie in cookie:
                    self.driver.add_cookie(thisCookie)
            case dict():
                self.driver.add_cookie(cookie)
            case _:
                return False
        return True


    def set_header(self, header):
        match header:
            case int() | float():
                return False
            case str():  # 字符串
                headers=[]
                header_strs = header.split(';')
                for headerStr in header_strs:
                    header=headerStr.split('=')
                    headers[len(headers)]={header[0]:header[1]}
                self.driver.execute_cdp_cmd("Network.setExtraHTTPHeaders", {"headers": headers})
            case list():
                for thisHeader in header:
                    self.driver.execute_cdp_cmd("Network.setExtraHTTPHeaders", {"headers": thisHeader})
            case dict():
                self.driver.execute_cdp_cmd("Network.setExtraHTTPHeaders", {"headers": header})
            case _:
                return False
        self.driver.execute_cdp_cmd("Network.enable", {})

    def get_html(self, url):
        self.driver.get(url)

    def screenshot(self):
        #网页截图
        return self.driver.get_screenshot_as_base64()

    def __del__(self):
        # 关闭浏览器
        self.driver.quit()
